#include <iostream>
#include <chrono>
#include <cmath>
#include <cstring>
#include <mpi.h>
#include "tensor_mat19.h"
#include "pcg.h"

// Our own implementations of CSR functions to avoid linking issues
void local_csr_spmv(const CsrMatrix &csr_matrix, double *x, double *b) {
    for (int i = 0; i < csr_matrix.rows; i++) {
        double sum = 0.0;
        for (int j = csr_matrix.row_off[i]; j < csr_matrix.row_off[i + 1]; j++) {
            sum += csr_matrix.data[j] * x[csr_matrix.cols[j]];
        }
        b[i] = sum;
    }
}

void local_free_csr_matrix(CsrMatrix &csr_matrix) {
    if (csr_matrix.row_off) delete[] csr_matrix.row_off;
    if (csr_matrix.cols) delete[] csr_matrix.cols;
    if (csr_matrix.data) delete[] csr_matrix.data;
    
    csr_matrix.row_off = nullptr;
    csr_matrix.cols = nullptr;
    csr_matrix.data = nullptr;
    csr_matrix.rows = 0;
    csr_matrix.data_size = 0;
}

// Local implementation of CG solver
void local_cg_solve(const CsrMatrix &A, const double *b, double *x, int max_iters, double tolerance, int *iter_out) {
    int n = A.rows;
    
    // Allocate memory for vectors
    double *r = new double[n]();  // residual
    double *p = new double[n]();  // search direction
    double *Ap = new double[n](); // A*p
    
    // Initialize r = b - A*x (assuming x = 0 initially)
    std::memcpy(r, b, n * sizeof(double));
    
    // Initialize p = r
    std::memcpy(p, r, n * sizeof(double));
    
    // Initial residual
    double r_norm = 0.0;
    for (int i = 0; i < n; i++) {
        r_norm += r[i] * r[i];
    }
    double initial_r_norm = r_norm;
    
    // Convergence check
    double tolerance_sq = tolerance * tolerance * initial_r_norm;
    
    // CG iterations
    int iter;
    for (iter = 0; iter < max_iters; iter++) {
        // Compute Ap = A*p
        local_csr_spmv(A, p, Ap);
        
        // Compute step size alpha = r*r / (p*Ap)
        double p_dot_Ap = 0.0;
        for (int i = 0; i < n; i++) {
            p_dot_Ap += p[i] * Ap[i];
        }
        double alpha = r_norm / p_dot_Ap;
        
        // Update solution: x = x + alpha*p
        for (int i = 0; i < n; i++) {
            x[i] += alpha * p[i];
        }
        
        // Update residual: r = r - alpha*Ap
        for (int i = 0; i < n; i++) {
            r[i] -= alpha * Ap[i];
        }
        
        // Compute new residual norm
        double new_r_norm = 0.0;
        for (int i = 0; i < n; i++) {
            new_r_norm += r[i] * r[i];
        }
        
        // Check for convergence
        if (new_r_norm < tolerance_sq) {
            iter++;
            break;
        }
        
        // Compute beta
        double beta = new_r_norm / r_norm;
        r_norm = new_r_norm;
        
        // Update search direction: p = r + beta*p
        for (int i = 0; i < n; i++) {
            p[i] = r[i] + beta * p[i];
        }
    }
    
    // Set output iteration count
    *iter_out = iter;
    
    // Clean up
    delete[] r;
    delete[] p;
    delete[] Ap;
}

// Function to initialize a vector with constant values
void init_constant_vector(float* vec, int size, float value) {
    for (int i = 0; i < size; i++) {
        vec[i] = value;
    }
}

// Function to convert float vector to double vector
void convert_float_to_double(const float* src, double* dst, int size) {
    for (int i = 0; i < size; i++) {
        dst[i] = static_cast<double>(src[i]);
    }
}

// Function to convert double vector to float vector
void convert_double_to_float(const double* src, float* dst, int size) {
    for (int i = 0; i < size; i++) {
        dst[i] = static_cast<float>(src[i]);
    }
}

// Function to print matrix statistics
void print_matrix_stats(const TensorMat19::TensorMatrix& mat) {
    const auto& grid = mat.grid;
    
    // Calculate some basic statistics
    float min_val = std::numeric_limits<float>::max();
    float max_val = std::numeric_limits<float>::lowest();
    float sum_val = 0.0f;
    
    // Sample some diagonal elements
    std::cout << "Sampling diagonal elements:" << std::endl;
    
    for (int k = 0; k < grid.kmt; k += grid.kmt/4) {
        for (int j = 0; j < grid.jmt; j += grid.jmt/4) {
            for (int i = 0; i < grid.imt; i += grid.imt/4) {
                for (int m = 1; m <= 19; m++) {
                    float val = mat.at(m, i, j, k);
                    min_val = std::min(min_val, val);
                    max_val = std::max(max_val, val);
                    sum_val += val;
                    
                    if (i == 0 && j == 0 && k == 0) {
                        std::cout << "  A(" << m << ", 0, 0, 0) = " << val << std::endl;
                    }
                }
            }
        }
    }
    
    std::cout << "Matrix statistics:" << std::endl;
    std::cout << "  Min value: " << min_val << std::endl;
    std::cout << "  Max value: " << max_val << std::endl;
    std::cout << "  Avg value: " << sum_val / (grid.total_inner_points * 19) << std::endl;
}

// Main test function
int main(int argc, char** argv) {
    // Initialize MPI
    int rank, size;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    if (rank == 0) {
        std::cout << "=== TensorMat19 Test Program ===" << std::endl;
    }
    
    // Parse command line arguments
    int resolution = 1;
    bool is_global_grid = true;
    
    for (int i = 1; i < argc; i++) {
        if (std::string(argv[i]) == "--resolution" && i + 1 < argc) {
            resolution = std::atoi(argv[i + 1]);
            i++;
        }
        else if (std::string(argv[i]) == "--local") {
            is_global_grid = false;
        }
    }
    
    // Create grid parameters
    TensorMat19::GridParams grid(resolution, is_global_grid);
    
    if (rank == 0) {
        grid.print_info();
    }
    
    // Create TensorMatrix
    TensorMat19::TensorMatrix tensor_mat(grid);
    
    // Initialize with random values
    if (rank == 0) {
        std::cout << "Initializing TensorMat19 with random values..." << std::endl;
    }
    
    auto start_time = std::chrono::high_resolution_clock::now();
    tensor_mat.init_random();
    auto end_time = std::chrono::high_resolution_clock::now();
    
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    if (rank == 0) {
        std::cout << "Initialization completed in " << duration << " ms" << std::endl;
        print_matrix_stats(tensor_mat);
    }
    
    // Create test vectors
    float* x = new float[grid.total_padded_points]();
    float* b = new float[grid.total_padded_points]();
    
    // Initialize b with all ones in the inner points
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                b[grid.padded_index(i, j, k)] = 1.0f;
            }
        }
    }
    
    // Apply boundary conditions to b
    TensorMat19::communicate_boundaries(b, grid);
    
    // Test matrix-vector multiplication
    if (rank == 0) {
        std::cout << "\nTesting matrix-vector multiplication..." << std::endl;
    }
    
    float* y = new float[grid.total_padded_points]();
    
    // Initialize x with zeros
    std::memset(x, 0, grid.total_padded_points * sizeof(float));
    
    // Perform y = A*x (should be zero)
    start_time = std::chrono::high_resolution_clock::now();
    tensor_mat.multiply(x, y);
    end_time = std::chrono::high_resolution_clock::now();
    
    duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    if (rank == 0) {
        std::cout << "Matrix-vector multiplication completed in " << duration << " ms" << std::endl;
        float y_norm = std::sqrt(TensorMat19::norm_squared(y, grid));
        std::cout << "||A*0|| = " << y_norm << " (should be close to 0)" << std::endl;
    }
    
    // Convert tensor matrix to CSR format
    if (rank == 0) {
        std::cout << "\nConverting TensorMat19 to CSR format..." << std::endl;
    }
    
    CsrMatrix csr_matrix;
    
    start_time = std::chrono::high_resolution_clock::now();
    tensor_mat.to_csr(csr_matrix);
    end_time = std::chrono::high_resolution_clock::now();
    
    duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    if (rank == 0) {
        std::cout << "Conversion completed in " << duration << " ms" << std::endl;
        std::cout << "CSR Matrix statistics:" << std::endl;
        std::cout << "  Rows: " << csr_matrix.rows << std::endl;
        std::cout << "  Non-zeros: " << csr_matrix.data_size << std::endl;
        std::cout << "  Avg non-zeros per row: " << static_cast<double>(csr_matrix.data_size) / csr_matrix.rows << std::endl;
    }
    
    // Test CSR matrix-vector multiplication
    if (rank == 0) {
        std::cout << "\nTesting CSR matrix-vector multiplication..." << std::endl;
    }
    
    // Convert float vectors to double for CSR operations
    double* x_double = new double[grid.total_inner_points]();
    double* b_double = new double[grid.total_inner_points]();
    double* y_double = new double[grid.total_inner_points]();
    
    // Extract the inner parts of b to b_double
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                int inner_idx = grid.inner_index(i, j, k);
                int padded_idx = grid.padded_index(i, j, k);
                b_double[inner_idx] = b[padded_idx];
            }
        }
    }
    
    // Test CSR SpMV
    start_time = std::chrono::high_resolution_clock::now();
    local_csr_spmv(csr_matrix, x_double, y_double);
    end_time = std::chrono::high_resolution_clock::now();
    
    duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    if (rank == 0) {
        std::cout << "CSR matrix-vector multiplication completed in " << duration << " ms" << std::endl;
        
        double y_norm = 0.0;
        for (int i = 0; i < grid.total_inner_points; i++) {
            y_norm += y_double[i] * y_double[i];
        }
        y_norm = std::sqrt(y_norm);
        
        std::cout << "||CSR*0|| = " << y_norm << " (should be close to 0)" << std::endl;
    }
    
    // Test CG solver
    if (rank == 0) {
        std::cout << "\nTesting CG solver with CSR matrix..." << std::endl;
    }
    
    // Initial guess (all zeros)
    double* solution = new double[grid.total_inner_points]();
    
    // Solve the system using CG
    int max_iterations = 1000;
    double tolerance = 1e-6;
    int iterations = 0;
    
    start_time = std::chrono::high_resolution_clock::now();
    local_cg_solve(csr_matrix, b_double, solution, max_iterations, tolerance, &iterations);
    end_time = std::chrono::high_resolution_clock::now();
    
    duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    if (rank == 0) {
        std::cout << "CG solver completed in " << iterations << " iterations (" << duration << " ms)" << std::endl;
        
        // Verify solution: A*x = b
        double* verify = new double[grid.total_inner_points]();
        local_csr_spmv(csr_matrix, solution, verify);
        
        double residual_norm = 0.0;
        double b_norm = 0.0;
        for (int i = 0; i < grid.total_inner_points; i++) {
            residual_norm += (verify[i] - b_double[i]) * (verify[i] - b_double[i]);
            b_norm += b_double[i] * b_double[i];
        }
        residual_norm = std::sqrt(residual_norm);
        b_norm = std::sqrt(b_norm);
        
        std::cout << "Relative residual: ||A*x - b|| / ||b|| = " << residual_norm / b_norm << std::endl;
        
        // Cleanup
        delete[] verify;
    }
    
    // Clean up resources
    delete[] x;
    delete[] b;
    delete[] y;
    delete[] x_double;
    delete[] b_double;
    delete[] y_double;
    delete[] solution;
    
    local_free_csr_matrix(csr_matrix);
    
    // Finalize MPI
    MPI_Finalize();
    
    return 0;
} 